#include "mainwindow.h"
#include "./ui_mainwindow.h"

DEFINE_string(pre_model, "/home/kirito/Desktop/segment-anything/sam_model/sam_preprocess.onnx", "Path to the preprocessing model");
DEFINE_string(sam_model, "/home/kirito/Desktop/segment-anything/sam_model/sam_vit_h_4b8939.onnx", "Path to the sam model");
DEFINE_string(pre_device, "cpu", "cpu or cuda:0(1,2,3...)");
DEFINE_string(sam_device, "cpu", "cpu or cuda:0(1,2,3...)");
DEFINE_bool(h, false, "Show help");

MainWindow::MainWindow(QWidget *parent)
    : QMainWindow(parent)
    , ui(new Ui::MainWindow)
{
    ui->setupUi(this);
}

MainWindow::~MainWindow()
{
    delete ui;
}

bool MainWindow::parseDeviceName(const std::string& name, Sam::Parameter::Provider& provider) {
    if (name == "cpu") {
        provider.deviceType = 0;
        return true;
    }
    if (name.substr(0, 5) == "cuda:") {
        provider.deviceType = 1;
        provider.gpuDeviceId = std::stoi(name.substr(5));
        return true;
    }
    return false;
}

void MainWindow::on_SegmentButton_clicked()
{
    std::cout << "Preprocess device: " << FLAGS_pre_device << "; Sam device: " << FLAGS_sam_device
              << std::endl;

    Sam::Parameter param(FLAGS_pre_model, FLAGS_sam_model, std::thread::hardware_concurrency());
    if (!parseDeviceName(FLAGS_pre_device, param.providers[0]) ||
        !parseDeviceName(FLAGS_sam_device, param.providers[1])) {
        std::cerr << "Unable to parse device name" << std::endl;
    }

    std::cout << "Loading model..." << std::endl;
    Sam sam(param);
    auto inputSize = sam.getInputSize();
    if (inputSize.empty()) {
        std::cout << "Sam initialization failed" << std::endl;
        return ;
    }

    cv::Mat image = cv::imread(this->_imageFilePath.toStdString(), -1);
    if (image.empty()) {
        std::cout << "Image loading failed" << std::endl;
        return ;
    }
    std::cout << "Resize image to " << inputSize << std::endl;
    cv::resize(image, image, inputSize);
    std::cout << "Loading image..." << std::endl;
    if (!sam.loadImage(image)) {
        std::cout << "Image loading failed" << std::endl;
        return ;
    }

    std::list<cv::Point3i> clickedPoints;
    cv::Point3i newClickedPoint(-1, 0, 0);
    cv::Rect roi;
    cv::Mat outImage = image.clone();

    cv::Size sampleSize = {image.cols / 20, image.rows / 20};

    std::cout << "Automatically generating masks with " << sampleSize.area()
              << " input points ..." << std::endl;

    auto mask = sam.autoSegment(
        sampleSize, [](double v) { std::cout << "\rProgress: " << int(v * 100) << "%\t"; },0.86);
    std::cout<<"auto segment"<< std::endl;

    //上色之前先保存被分割物体//
    QFileInfo fileInfo(this->_imageFilePath);
    qDebug() << fileInfo;
    QString fileName = fileInfo.completeBaseName();
    qDebug() << fileName;
    QListWidgetItem* item = ui->OutImageList->findItems(fileInfo.fileName(), Qt::MatchExactly).value(0);
    saveSegmentedObjects(mask,image,fileName.toStdString(),item);

    const double overlayFactor = 0.5;
    const int maxMaskValue = 255 * (1 - overlayFactor);
    outImage = cv::Mat::zeros(image.size(), CV_8UC3);

    static std::map<int, cv::Vec3b> colors;

    //outImage色彩输出

    for (int i = 0; i < image.rows; i++) {
        for (int j = 0; j < image.cols; j++) {
            auto value = (int)mask.at<double>(i, j);
            if (value <= 0) {
                continue;
            }

            auto it = colors.find(value);
            if (it == colors.end()) {
                colors.insert(it, {value, cv::Vec3b(rand() % maxMaskValue, rand() % maxMaskValue,
                                                    rand() % maxMaskValue)});
            }

            outImage.at<cv::Vec3b>(i, j) = it->second + image.at<cv::Vec3b>(i, j) * overlayFactor;
        }
    }

    // outImage转换成ui->label上面
    QImage qimage = QImage((const uchar*)outImage.data,
                           outImage.cols,
                           outImage.rows,
                           outImage.step,
                           QImage::Format_RGB888);

    // 转换后的QImage可以转换为QPixmap，然后设置到QLabel上
    QPixmap qpixmap = QPixmap::fromImage(qimage);

    // 假设您有一个名为 ui->label 的QLabel对象
    ui->ImageDisplay->setPixmap(qpixmap.scaled(ui->ImageDisplay->size(), Qt::KeepAspectRatio, Qt::SmoothTransformation));
}


void MainWindow::saveSegmentedObjects(const cv::Mat& outImage, const cv::Mat& originalImage, const std::string& saveDir, QListWidgetItem* originalImageItem)
{
    if (outImage.empty() || originalImage.empty()) return;

    // 获取当前应用程序的工作目录
    QString currentDir = QCoreApplication::applicationDirPath();

    // 构造输出目录路径
    QString outputDir = QDir(currentDir).filePath("output");

    // 确保输出目录存在
    QDir().mkpath(outputDir);

    // 使用 qDebug() 输出带 "output" 目录的完整路径
    qDebug() << "Output directory set to:" << outputDir;

    std::vector<std::string> savedImagePaths;
    std::set<int> uniqueIndices;

    for (int i = 0; i < outImage.rows; ++i) {
        for (int j = 0; j < outImage.cols; ++j) {
            int index = static_cast<int>(outImage.at<double>(i, j));
            if (index > 0) {
                uniqueIndices.insert(index);
            }
        }
    }

    for (int index : uniqueIndices) {
        cv::Mat objectMask = cv::Mat::zeros(outImage.size(), CV_8UC1);
        for (int i = 0; i < outImage.rows; ++i) {
            for (int j = 0; j < outImage.cols; ++j) {
                if (static_cast<int>(outImage.at<double>(i, j)) == index) {
                    objectMask.at<uchar>(i, j) = 255;
                }
            }
        }

        std::vector<std::vector<cv::Point>> contours;
        cv::findContours(objectMask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

        if (contours.empty()) continue;

        cv::Mat croppedObject = cv::Mat::zeros(objectMask.size(), CV_8UC4);
        cv::Mat objectMaskColor;
        cv::cvtColor(objectMask, objectMaskColor, cv::COLOR_GRAY2BGR);
        originalImage.copyTo(croppedObject, objectMaskColor);

        if (cv::countNonZero(objectMask) < 256)
            continue;

        std::string filename = outputDir.toStdString() + "/object_" + std::to_string(index) + ".png";
        cv::imwrite(filename, croppedObject);
        std::cout << "Saved object " << index << " to " << filename << std::endl;

        QFileInfo fileInfo(QString::fromStdString(filename));
        QString imageName = fileInfo.fileName();
        qDebug()<<imageName<<" "<<QString::fromStdString(filename);
        imageMap.insert(imageName, QString::fromStdString(filename));
        savedImagePaths.push_back(filename); // 保存路径用于后续添加到 QListWidget
    }

    // 更新 QListWidget
    addImagesToListWidget(savedImagePaths, originalImageItem);
}


void MainWindow::addImagesToListWidget(const std::vector<std::string>& imagePaths, QListWidgetItem* originalImageItem)
{
    // 检查 QListWidget 是否存在
    if (!ui->OutImageList) {
        std::cerr << "QListWidget not found!" << std::endl;
        return;
    }

    // 遍历图像路径并将图像添加到 QListWidget
    for (const auto& path : imagePaths) {
        // 创建新的 QListWidgetItem
        QListWidgetItem* newItem = new QListWidgetItem;

        // 将图像路径转换为 QImage
        QImage image(QString::fromStdString(path));
        if (image.isNull()) {
            std::cerr << "Failed to load image: " << path << std::endl;
            continue;
        }

        // 将 QImage 转换为 QPixmap
        QPixmap pixmap = QPixmap::fromImage(image);

        // 设置 QListWidgetItem 的图标
        newItem->setIcon(QIcon(pixmap));


        QFileInfo fileInfo(QString::fromStdString(path));
        QString fileName = fileInfo.fileName();
        // 设置 QListWidgetItem 的文本
        newItem->setText(fileName);

        // 将新的 QListWidgetItem 添加到 QListWidget 中
        ui->OutImageList->addItem(newItem);
    }

    if (originalImageItem) {
        ui->OutImageList->setCurrentItem(originalImageItem);
    }
}


void MainWindow::on_pushButton_clicked()
{
    QString filter = "Image files (*.png *.jpg *.jpeg)";
    this->_imageFilePath = QFileDialog::getOpenFileName(this,
                                                         tr("Select Image"),
                                                         QDir::homePath(),
                                                         filter);

    if (this->_imageFilePath.isEmpty()) {
        QMessageBox::information(this, tr("Information"), tr("No image selected."));
    } else {
        // 使用 QFileInfo 提取文件名
        QFileInfo fileInfo(this->_imageFilePath);
        QString fileName = fileInfo.fileName(); // 获取文件名
        if (ui->OutImageList->findItems(fileName, Qt::MatchExactly).isEmpty())
        {
            imageMap.insert(fileName, this->_imageFilePath);

            // 将只含文件名的项添加到 QListWidget 中
            QListWidgetItem *item = new QListWidgetItem(QIcon(), fileName);
            ui->OutImageList->addItem(item);
        }
        // 找到重复项，提示用户
        else QMessageBox::information(this, tr("Information"), tr("图片已经在列表中。"));
    }
}


void MainWindow::on_OutImageList_itemClicked(QListWidgetItem *item)
{
    // 从项中获取图片名称
    QString imageName = item->text();

    // 使用图片名作为键在 QMap 中查找对应的路径
    this->_imageFilePath = imageMap.value(imageName);

    // 检查路径是否存在
    if (!this->_imageFilePath.isEmpty()) {
        // 使用找到的路径加载图片
        QImage image(this->_imageFilePath);

        // 如果图片加载失败，则显示错误消息
        if (image.isNull()) {
            QMessageBox::warning(this, tr("Warning"), tr("Cannot load the image."));
            return;
        }

        // 将图像转换为 QPixmap，QLabel 可以直接显示 QPixmap
        QPixmap pixmap = QPixmap::fromImage(image);

        // 调整 QPixmap 的尺寸以适应 QLabel（如果需要）
        pixmap = pixmap.scaled(ui->ImageDisplay->size(), Qt::KeepAspectRatio, Qt::SmoothTransformation);

        // 将 QPixmap 设置到 QLabel 上
        ui->ImageDisplay->setPixmap(pixmap);
    } else {
        // 如果在 QMap 中没有找到对应的路径，显示错误消息
        QMessageBox::warning(this, tr("Warning"), tr("Image path not found."));
    }
}
